import numpy as np
import mxnet as mx
import mxnet.gluon as gluon
import mxnet.gluon.nn as nn
import gym


D = 80 * 80
H = 200
RENDER = True


def prepro(I):
    """ prepro 210x160x3 uint8 frame into 6400 (80x80) 1D float vector """
    I = I[35:195]  # crop
    I = I[::2, ::2, 0]  # downsample by factor of 2
    I[I == 144] = 0  # erase background (background type 1)
    I[I == 109] = 0  # erase background (background type 2)
    I[I != 0] = 1  # everything else (paddles, ball) just set to 1
    return I.astype(np.float).ravel()


class PG(nn.Block):
    def __init__(self):
        super(PG, self).__init__()
        self.l1 = nn.Dense(H)
        self.l2 = nn.Dense(1)

    def forward(self):
        pass


def train():
    env = gym.make('Pong-v0')
    observation = env.reset()
    prev_x = None
    running_reward = None
    reward_sum = 0
    episode_number = 0

    while True:
        if RENDER:
            env.render()

            # preprocess the observation of the environment
            cur_x = prepro(observation)
            x = cur_x - prev_x if prev_x is not None else 0
            prev_x = cur_x

            observation, reward, done, info = env.step(1)


if __name__ == '__main__':
    train()
